import numpy as np
import pandas as pd
import cvxpy as cp

# 区域人口占比
area_target_dict = {
    1: 0.067797428,  # 东湖区	    1
    2: 0.078006431,  # 西湖区	    2
    3: 0.056125402,  # 青云谱区	3
    4: 0.106527331,  # 青山湖区	4
    5: 0.106784566,  # 新建区	    5
    6: 0.089356913,  # 红谷滩区	6
    7: 0.074485531,  # 经开区	    7
    8: 0.071881029,  # 高新区	    8
    9: 0.017893891,  # 湾里区	    9
    10: 0.190803859,  # 南昌县	    10
    11: 0.040610932,  # 安义县	    11
    12: 0.099726688  # 进贤县	    12
}

# 年龄段人口占比
age_target_dict = [
    {'lb': 0, 'ub': 14, 'num': 0.1744},
    {'lb': 15, 'ub': 59, 'num': 0.6759},
    {'lb': 60, 'ub': 120, 'num': 0.1497}
]

# 教育人口占比
education_target_dict = [
    {'lb': 1, 'ub': 2, 'num': 0.220386594}, # 小学以下 	1 2
    {'lb': 3, 'ub': 3, 'num': 0.324917923}, # 初中	3
    {'lb': 4, 'ub': 4, 'num': 0.179901333}, # 高中	4
    {'lb': 5, 'ub': 7, 'num': 0.27479415}, # 大学及以上	5 6 7
]

# 读取问卷
all_item = pd.read_excel("./data.xlsx", sheet_name=0)

# item_var_list = [] # list of variable
area_list = []
age_list = []
education_list = []

# generate variables and parameters
item_num = len(all_item)
x = cp.Variable(item_num, nonneg=True)
# for key, item in all_item.iloc[:400].iterrows():
for key, item in all_item.iterrows():
    try:
        item_id = item['ID']
        area_list.append(item['6居住地'])
        age_list.append(item['2年龄'])
        education_list.append(item['4文化程度'])
    except Exception as e:
        print(f'{item_id} goes wrong!')
        continue

# 权重之和
weight_sum = 10000

def loss_fn(X, Y, beta, weight_sum):
    return cp.norm2(X @ beta - weight_sum * Y)

def loss_fn_real(X, Y, beta, weight_sum):
    return np.sqrt((np.multiply(X, beta).sum() - weight_sum * Y) ** 2)

# 为每个指标准备矩阵
expr_list = []

# expr 区域人口
for i in range(12):
    beta = np.zeros(item_num)
    for k in range(item_num):  # TODO: 筛选值等于1的样本
        if area_list[k] == i + 1:
            beta[k] = 1
    expr = loss_fn(beta, area_target_dict[i + 1], x, weight_sum)
    expr_list.append(expr)

# expr 年龄段人口
for i in range(3):
    beta = np.zeros(item_num)
    for k in range(item_num):
        if age_target_dict[i]['lb'] <= age_list[k] <= age_target_dict[i]['ub']:
            beta[k] = 1
    expr = loss_fn(beta, age_target_dict[i]['num'], x, weight_sum)
    expr_list.append(expr)

# 教育人口占比
for i in range(4):
    beta = np.zeros(item_num)
    for k in range(item_num):
        if education_target_dict[i]['lb'] <= education_list[k] <= education_target_dict[i]['ub']:
            beta[k] = 1
    expr = loss_fn(beta, education_target_dict[i]['num'], x, weight_sum)
    expr_list.append(expr)


problem = cp.Problem(cp.Minimize(sum(expr_list)), [x @ np.ones(item_num) <= weight_sum, x @ np.ones(item_num) >= weight_sum])

problem.solve(verbose=True)

# Print result.
print("\nThe optimal value is", problem.value)
real_weight_sum = x.value.sum()
print(f'Check Sum of Weight: {real_weight_sum}')
obj_sum = 0.0
# Check Error
for i in range(12):
    beta = np.zeros(item_num)
    for k in range(item_num):  # TODO: 筛选值等于1的样本
        if area_list[k] == i + 1:
            beta[k] = 1
    obj = loss_fn_real(beta, area_target_dict[i+1], x.value, real_weight_sum)
    obj_sum += float(obj)
for i in range(3):
    beta = np.zeros(item_num)
    for k in range(item_num):
        if age_target_dict[i]['lb'] <= age_list[k] <= age_target_dict[i]['ub']:
            beta[k] = 1
    obj = loss_fn_real(beta, age_target_dict[i]['num'], x.value, real_weight_sum)
    obj_sum += float(obj)
for i in range(4):
    beta = np.zeros(item_num)
    for k in range(item_num):
        if education_target_dict[i]['lb'] <= education_list[k] <= education_target_dict[i]['ub']:
            beta[k] = 1
    obj = loss_fn_real(beta, education_target_dict[i]['num'], x.value, real_weight_sum)
    obj_sum += float(obj)
print(f'Error Value: {str(obj_sum)}')

# Calculate Origin Error
origin = np.ones(item_num)
for i in range(12):
    beta = np.zeros(item_num)
    for k in range(item_num):  # TODO: 筛选值等于1的样本
        if area_list[k] == i + 1:
            beta[k] = 1
    obj = loss_fn_real(beta, area_target_dict[i+1], origin, item_num)
    obj_sum += float(obj)
for i in range(3):
    beta = np.zeros(item_num)
    for k in range(item_num):
        if age_target_dict[i]['lb'] <= age_list[k] <= age_target_dict[i]['ub']:
            beta[k] = 1
    obj = loss_fn_real(beta, age_target_dict[i]['num'], origin, item_num)
    obj_sum += float(obj)
for i in range(4):
    beta = np.zeros(item_num)
    for k in range(item_num):
        if education_target_dict[i]['lb'] <= education_list[k] <= education_target_dict[i]['ub']:
            beta[k] = 1
    obj = loss_fn_real(beta, education_target_dict[i]['num'], origin, item_num)
    obj_sum += float(obj)
print(f'Origin Error Value: {str(obj_sum)}')

round_x = np.round(x.value * 10, decimals=0)
round_weight_sum = round_x.sum()
# Calculate Rounded Error
for i in range(12):
    beta = np.zeros(item_num)
    for k in range(item_num):  # TODO: 筛选值等于1的样本
        if area_list[k] == i + 1:
            beta[k] = 1
    obj = loss_fn_real(beta, area_target_dict[i+1], round_x, round_weight_sum)
    obj_sum += float(obj)
for i in range(3):
    beta = np.zeros(item_num)
    for k in range(item_num):
        if age_target_dict[i]['lb'] <= age_list[k] <= age_target_dict[i]['ub']:
            beta[k] = 1
    obj = loss_fn_real(beta, age_target_dict[i]['num'], round_x, round_weight_sum)
    obj_sum += float(obj)
for i in range(4):
    beta = np.zeros(item_num)
    for k in range(item_num):
        if education_target_dict[i]['lb'] <= education_list[k] <= education_target_dict[i]['ub']:
            beta[k] = 1
    obj = loss_fn_real(beta, education_target_dict[i]['num'], round_x, round_weight_sum)
    obj_sum += float(obj)
print(f'Rounded Error Value: {str(obj_sum)}')

# Export Excel File
df = pd.DataFrame(x.value)
df.to_excel('output.xlsx', index=False)